"""
The netpop.manager package has code for the ExperimentManager, a class that automates job submission and execution,
and the ExperimentRunner, a process for running jobs.

Perhaps I should keep track of the host that checks out a job for processing.
"""
import os, sys, shelve, collections, xmlrpclib, time, datetime
import cPickle as pickle
from SimpleXMLRPCServer import SimpleXMLRPCServer
try:
	from elixir import *
except:
	print 'Elixir is not available on this system.'
	
import sqlite3, zlib

from bourbon.util import EM_PORT, exp_from_data, data_from_exp, load_defaults

#EM_PORT = 8000
#bourbon manager --db=mysql://bourbon:bourbon@localhost/bourbon --host=leviathan.cs.rice.edu --port=8000
#bourbon manager --db=mysql://bourbon:bourbon@localhost/bourbon --host=localhost --port=8000
"""
class Result(Entity):
	using_options(shortnames=True)
	
	name = Field(Unicode(100))
	datatype = Field(Unicode(20))
	value = Field(Float)
	text = Field(Unicode(255))
	num = Field(Integer)
	data = Field(Binary)
	description = Field(UnicodeText)
	submitted = Field(DateTime,default=datetime.datetime.now)
	job = ManyToOne('Job')
	
	def get_value(self):
		if self.datatype == 'int':
			return self.num
		elif self.datatype == 'float':
			return self.value
		elif self.datatype == 'str':
			return self.text
		elif self.datatype == 'object':
			return pickle.loads(self.data)
			
		return None
			
	def __repr__(self):
		return "%s=%s"%(self.name, str(self.get_value()))
		
class Parameter(Entity):
	using_options(shortnames=True)
	name = Field(Unicode(100))
	datatype = Field(Unicode(20))
	value = Field(Float)
	text = Field(Unicode(255))
	num = Field(Integer)
	data = Field(Binary)
	description = Field(UnicodeText)
	job = ManyToOne('Job')
	
	def get_value(self):
		if self.datatype == 'int':
			return self.num
		elif self.datatype == 'float':
			return self.value
		elif self.datatype == 'str':
			return self.text
		elif self.datatype == 'object':
			return pickle.loads(self.data)
			
		return None
			
	def __repr__(self):
		return "%s=%s"%(self.name, str(self.get_value()))
		
# Change this to Run
class Job(Entity):
	using_options(shortnames=True)
	submitted = Field(DateTime,default=datetime.datetime.now)
	name = Field(Unicode(30))
	status = Field(Unicode(10),default=JobStatus.PENDING)
	completed = Field(DateTime)
	duration = Field(Float, default=0.0)
	description = Field(UnicodeText)
	data = OneToMany('Data')
	design = ManyToOne('Design')
	parameters = OneToMany('Parameter')
	results = OneToMany('Result')
	
	def add_param(self, name, value):
		param = Parameter(name=unicode(name), job=self)
		self.set_value(param, value)
		
	def add_result(self, name, value):
		result = Result(name=unicode(name), job=self)
		self.set_value(result, value)
		
	def set_value(self, obj, value):
		if isinstance(value, int):
			obj.num = value
			obj.value = value
			obj.datatype = u'int'
		elif isinstance(value, float):
			obj.value = value
			obj.datatype = u'float'
		elif isinstance(value, str):
			obj.text = value
			obj.datatype = u'str'
		else:
			obj.text = str(value)
			obj.data = pickle.dumps(value)
			obj.description = str(value.__class__)
			obj.datatype = u'object'

	def get_param_value(self, name):
		for param in self.parameters:
			if param.name.endswith(name):
				return param.get_value()
		return None
		
	def get_result_value(self, name):
		for result in self.results:
			if result.name.endswith(name):
				return result.get_value()
		return None
	
# Change this to Job
class Design(Entity):
	using_options(shortnames=True)
	saved = Field(DateTime, default=datetime.datetime.now)
	name = Field(Unicode(30))
	data = Field(Binary)
	runs = OneToMany('Job')
	
	def load(self):
		return pickle.loads(self.data)
	
class Data(Entity):
	using_options(shortnames=True)
	
	saved = Field(DateTime,default=datetime.datetime.now)
	progress = Field(Unicode(10))
	data = Field(Binary)
	run = ManyToOne('Job')
	description = Field(UnicodeText)
	
	def load(self):
		return pickle.loads(zlib.decompress(self.data))
"""

class NamedValue(Entity):
	using_options(shortnames=True)
	name = Field(String(100))
	datatype = Field(String(20))
	value_number = Field(Float)
	value_text = Field(String(255))
	value_int = Field(Integer)
	value = Field(PickleType)
	description = Field(String(255))
	
	def __init__(self, name, value, description=None):
		self.name = name
		self.description = description
		self.set_value(value)
		
	def set_value(self, value):
		if isinstance(value, int):
			self.value_int = value
			self.value_number = value
			self.datatype = 'int'
		elif isinstance(value, float):
			self.value_number = value
			self.datatype = 'float'
		elif isinstance(value, str):
			self.value_text = value
			self.datatype = 'str'
		else:
			self.value_text = str(value)
			self.datatype = 'object'
		self.value = value
	
	def get_value(self):
		return self.value
			
	def __repr__(self):
		return "%s=%s"%(self.name, str(self.get_value()))
		
class Parameter(NamedValue):
	job = ManyToOne('Job')
	
class Result(NamedValue):
	run = ManyToOne('Run')
	
class Action(Entity):
	using_options(shortnames=True)
	action_type = Field(String(20))
	description = Field(String(255))
	date = Field(DateTime, default=datetime.datetime.now)
	data = Field(PickleType)
	run = ManyToOne('Run')
	
class Run(Entity):
	using_options(shortnames=True)
	submitted = Field(DateTime,default=datetime.datetime.now)
	name = Field(String(100))
	status = Field(String(20), default="PENDING")
	completed = Field(DateTime)
	description = Field(Text)
	actions = OneToMany('Action')
	job = ManyToOne('Job')
	results = OneToMany('Result')
	
	def get_param_value(self, name):
		for param in self.parameters:
			if param.name.endswith(name):
				return param.get_value()
		return None
		
	def get_result_value(self, name):
		for result in self.results:
			if result.name.endswith(name):
				return result.get_value()
		return None
		
	def data(self):
		"""Retrieves all data associated with this Run."""
		data = {}
		for result in self.results:
			data[result.name] = result.get_value()
		for param in self.job.params:
			data[param.name] = param.get_value()
			
		return data
	
class Job(Entity):
	using_options(shortnames=True)
	saved = Field(DateTime, default=datetime.datetime.now)
	name = Field(String(30))
	data = Field(PickleType)
	description = Field(Text)
	params = OneToMany('Parameter')
	runs = OneToMany('Run')
	
	
class JobManager:
	def __init__(self, db_address=None):
		self._db_address = db_address
		self.__inner = None
		
	def inner(self):
		#return InnerJobManager(self._db_address)
		if self.__inner is None: self.__inner = InnerJobManager(self._db_address)
		return self.__inner
		
	def clear_jobs(self):
		self.inner().clear_jobs()

	def submit_job(self, job, name="NoName", times=1, description=None):
		"""Submits the job to the manager.  A unique ID will be provided for the job.
			
			The ExperimentManager will save the job and put it into the job queue.
		"""
		return self.inner().submit_job(job, name, times, description)
	
	def check(self):
		return 1
	
	def submit_results_binary(self, run_id, data):
		return self.inner().submit_results_binary(run_id, data)

	def submit_results(self, run_id, results):
		"""Saves the results of this job to the output folder."""
		return self.inner().submit_results(run_id, results)

	def next_work(self):
		return self.inner().next_work()		

	def checkout_success(self, run_id, success):
		return self.inner().checkout_success(run_id, success)		

	def get_next_work(self):
		return self.inner().get_next_work()
	
	
	def get_work_data(self, job_id,rpc=False):
		return self.inner().get_work_data(job_id, rpc)
		
		
	def start_server(self, host='localhost',port=EM_PORT):
		addr = (host, port)
		print 'STARTING on %s:%d'%(host, port)
		server = SimpleXMLRPCServer(addr, allow_none=True)
		server.register_instance(self)
		server.serve_forever()

	def __repr__(self):
		data = "JobManager\n"
		data += "\tDatabase=%s\n"%str(metadata.bind)
		data += "\tExperiments=%d"%len(self.inner().runs())
		return data
		
class InnerJobManager:
	
	def __init__(self, db_address=None):
		if db_address is None:
			self.db_address = self.default_db()
		else:
			self.db_address = db_address
		self.rebind()
		
	def rebind(self):
		metadata.bind = self.db_address
		setup_all(True)
		
	def default_db(self):
		# check to see if there is a .bourbon file
		data = load_defaults()
		if 'db_address' in data:
			return data['db_address']
		
		path = os.path.join(os.getenv('HOME'), 'bourbon_db.sqlite')
		if not os.path.exists(path):
			# create the sqlite file
			print 'Creating sqlite database at ', path
			conn = sqlite3.connect(path)
			conn.close()
		
		return "sqlite:///%s"%path
		
		
	def clear_jobs(self):
		for run in self.query_runs():
			run.status = 'CANCELLED'
		session.commit()
		
	def submit_job(self, job, name="NoName", times=1, description=None):
		"""Submits the job to the manager.  A unique ID will be provided for the job.
		
		The ExperimentManager will save the job and put it into the job queue.
		"""
		einit = Job(data=job, name=name, description=description)
		eruns = [ Run(name=name, job=einit, status='PENDING',actions=[Action(action_type='SUBMITTED')]) for i in range(times) ]
		# Store parameters for each job
		
		for k,v in job.params().iteritems():
			einit.params.append(Parameter(k,v))
				
		session.commit()
		result = [er.id for er in eruns]
		if times == 1:
			return result[0]
		return result
		
	def submit_results_binary(self, run_id, data):
		self.submit_results(run_id, data.data)
		
	def submit_results(self, run_id, results):
		"""Saves the results of this job to the output folder."""
		print 'Saving results for [%d]...'%run_id,
		erun = self.fetch(run_id)
		results = pickle.loads(results)
		#if isinstance(data, xmlrpclib.Binary): data = data.data
		
		for k,v in results.iteritems():
			erun.results.append(Result(k,v))
		
		erun.status = 'FINISHED'
		erun.completed = datetime.datetime.now()
		erun.actions.append(Action(action_type='FINISHED'))
		# Store results
		
		session.commit()
		
		try:
			session.commit()
		except:
			print 'ERROR!'
		else:
			print 'DONE.'
		return 1
		
		
	def fetch(self, eid):
		return Run.query.filter_by(id=eid).one()
		
	def query_runs(self):
		#self.rebind()
		q = Run.query.filter_by(status='PENDING')
		q.order_by(Run.submitted)
		return q
		
	def runs(self):
		return self.query_runs().all()
	
	def all_runs(self):
		q = Run.query.order_by(Run.submitted)
		return q.all()
	
	def has_work(self):
		"""Returns True if there are more jobs in the queue."""
		return self.query_runs().count() > 0
	
	def next_work(self):
		
		if self.has_work():
			erun = self.query_runs().first()
			
			erun.status = 'CHECKOUT'
			#erun.actions.append(Action( action_type='CHECKOUT'))
			
			session.commit()
			return erun
			
		return None
	
	def checkout_success(self, run_id, success):
		erun = self.fetch(run_id)
		if success:
			erun.status = 'RUNNING'
			erun.actions.append(Action( action_type='CHECKOUT'))
		else:
			erun.status = 'PENDING'
			print 'ERROR CHECKING OUT [%d]'%run_id
		session.commit()
		return 0
		
	def get_next_work(self):
		#print 'Getting next work item...'
		#print self.query_runs().count()
		exp = self.next_work()
		if exp is None: return -1
		
		return exp.id
	
	def get_work_data(self, job_id,rpc=False):
		exp = self.fetch(job_id)
		data = exp.job.data
		if rpc:
			data = xmlrpclib.Binary(data_from_exp(data))
		else:
			data = data
		return data
		
	"""
	def start_server(self, host='localhost',port=EM_PORT):
		addr = (host, port)
		print 'STARTING on %s:%d'%(host, port)
		server = SimpleXMLRPCServer(addr, allow_none=True)
		server.register_instance(self)
		
		server.serve_forever()
	"""
	
	def __repr__(self):
		data = "JobManager\n"
		data += "\tDatabase=%s\n"%str(metadata.bind)
		data += "\tExperiments=%d"%len(self.runs())
		return data
	

			
			
	
	

